"""
基础记忆模块

定义记忆系统的抽象基类和通用接口
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional, Union

from langchain_core.language_models import BaseLanguageModel
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage


@dataclass
class MemoryItem:
    """记忆项"""
    id: str
    content: str
    message_type: str
    timestamp: datetime
    metadata: Dict[str, Any] = field(default_factory=dict)
    importance: float = 0.5  # 重要性评分 (0-1)
    tags: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'id': self.id,
            'content': self.content,
            'message_type': self.message_type,
            'timestamp': self.timestamp.isoformat(),
            'metadata': self.metadata,
            'importance': self.importance,
            'tags': self.tags
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'MemoryItem':
        """从字典创建"""
        return cls(
            id=data['id'],
            content=data['content'],
            message_type=data['message_type'],
            timestamp=datetime.fromisoformat(data['timestamp']),
            metadata=data.get('metadata', {}),
            importance=data.get('importance', 0.5),
            tags=data.get('tags', [])
        )


@dataclass
class MemorySearchResult:
    """记忆搜索结果"""
    item: MemoryItem
    relevance_score: float
    match_reason: str
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'item': self.item.to_dict(),
            'relevance_score': self.relevance_score,
            'match_reason': self.match_reason
        }


class BaseMemory(ABC):
    """
    记忆系统抽象基类
    
    定义所有记忆系统必须实现的接口
    """
    
    def __init__(self, agent_id: str, llm: Optional[BaseLanguageModel] = None):
        """
        初始化记忆系统
        
        Args:
            agent_id: 智能体ID
            llm: 大语言模型（可选）
        """
        self.agent_id = agent_id
        self.llm = llm
        self.created_at = datetime.now()
        self.updated_at = datetime.now()
        self._message_count = 0
        self._total_tokens = 0
    
    @abstractmethod
    def add_message(self, message: Union[str, BaseMessage], 
                   message_type: str = "human", 
                   metadata: Optional[Dict[str, Any]] = None) -> str:
        """
        添加消息到记忆
        
        Args:
            message: 消息内容
            message_type: 消息类型
            metadata: 元数据
            
        Returns:
            消息ID
        """
        pass
    
    @abstractmethod
    def get_messages(self, limit: Optional[int] = None, 
                    message_type: Optional[str] = None) -> List[MemoryItem]:
        """
        获取消息列表
        
        Args:
            limit: 限制数量
            message_type: 消息类型过滤
            
        Returns:
            消息列表
        """
        pass
    
    @abstractmethod
    def search_messages(self, query: str, limit: int = 10, 
                       threshold: float = 0.7) -> List[MemorySearchResult]:
        """
        搜索消息
        
        Args:
            query: 搜索查询
            limit: 限制数量
            threshold: 相关性阈值
            
        Returns:
            搜索结果列表
        """
        pass
    
    @abstractmethod
    def clear_memory(self) -> None:
        """清空记忆"""
        pass
    
    @abstractmethod
    def get_memory_stats(self) -> Dict[str, Any]:
        """
        获取记忆统计信息
        
        Returns:
            统计信息字典
        """
        pass
    
    def add_conversation_turn(self, human_message: str, ai_response: str,
                             metadata: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
        """
        添加对话轮次
        
        Args:
            human_message: 人类消息
            ai_response: AI回复
            metadata: 元数据
            
        Returns:
            消息ID字典
        """
        human_id = self.add_message(human_message, "human", metadata)
        ai_id = self.add_message(ai_response, "ai", metadata)
        
        return {
            "human_message_id": human_id,
            "ai_message_id": ai_id
        }
    
    def get_conversation_history(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        获取对话历史
        
        Args:
            limit: 限制数量
            
        Returns:
            对话历史列表
        """
        messages = self.get_messages(limit)
        conversation = []
        
        for message in messages:
            conversation.append({
                "role": message.message_type,
                "content": message.content,
                "timestamp": message.timestamp,
                "metadata": message.metadata
            })
        
        return conversation
    
    def summarize_conversation(self, max_length: int = 1000) -> str:
        """
        总结对话内容
        
        Args:
            max_length: 最大长度
            
        Returns:
            对话总结
        """
        if not self.llm:
            return "无法生成总结：未提供语言模型"
        
        messages = self.get_messages()
        if not messages:
            return "没有对话历史"
        
        # 构建总结提示
        conversation_text = "\n".join([
            f"{msg.message_type}: {msg.content}" 
            for msg in messages[-10:]  # 只取最近10条消息
        ])
        
        prompt = f"""
请总结以下对话内容，保持简洁明了：

{conversation_text}

总结：
"""
        
        try:
            response = self.llm.invoke(prompt)
            summary = response.content if hasattr(response, 'content') else str(response)
            
            # 限制长度
            if len(summary) > max_length:
                summary = summary[:max_length] + "..."
            
            return summary
        except Exception as e:
            return f"总结生成失败: {str(e)}"
    
    def extract_key_information(self, query: str) -> List[MemorySearchResult]:
        """
        提取关键信息
        
        Args:
            query: 查询内容
            
        Returns:
            相关信息列表
        """
        return self.search_messages(query, limit=5, threshold=0.6)
    
    def get_recent_context(self, limit: int = 5) -> List[MemoryItem]:
        """
        获取最近的上下文
        
        Args:
            limit: 限制数量
            
        Returns:
            最近的记忆项列表
        """
        return self.get_messages(limit=limit)
    
    def update_message_importance(self, message_id: str, importance: float) -> bool:
        """
        更新消息重要性
        
        Args:
            message_id: 消息ID
            importance: 重要性评分 (0-1)
            
        Returns:
            是否更新成功
        """
        # 默认实现，子类可以重写
        return False
    
    def add_tags_to_message(self, message_id: str, tags: List[str]) -> bool:
        """
        为消息添加标签
        
        Args:
            message_id: 消息ID
            tags: 标签列表
            
        Returns:
            是否添加成功
        """
        # 默认实现，子类可以重写
        return False
    
    def get_messages_by_tags(self, tags: List[str]) -> List[MemoryItem]:
        """
        根据标签获取消息
        
        Args:
            tags: 标签列表
            
        Returns:
            消息列表
        """
        # 默认实现，子类可以重写
        return []
    
    def _generate_message_id(self) -> str:
        """生成消息ID"""
        self._message_count += 1
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return f"{self.agent_id}_{timestamp}_{self._message_count}"
    
    def _update_stats(self, message: str) -> None:
        """更新统计信息"""
        self._message_count += 1
        self._total_tokens += len(message.split())  # 简单的词数统计
        self.updated_at = datetime.now()
    
    def _normalize_message(self, message: Union[str, BaseMessage]) -> str:
        """标准化消息格式"""
        if isinstance(message, BaseMessage):
            return message.content
        return str(message)
    
    def _get_message_type(self, message: Union[str, BaseMessage]) -> str:
        """获取消息类型"""
        if isinstance(message, HumanMessage):
            return "human"
        elif isinstance(message, AIMessage):
            return "ai"
        elif isinstance(message, SystemMessage):
            return "system"
        else:
            return "unknown"
    
    def __len__(self) -> int:
        """返回记忆项数量"""
        return self._message_count
    
    def __str__(self) -> str:
        return f"BaseMemory(agent_id='{self.agent_id}', messages={self._message_count})"
    
    def __repr__(self) -> str:
        return f"<BaseMemory(agent_id='{self.agent_id}', messages={self._message_count})>"


class MemoryManager:
    """记忆管理器"""
    
    def __init__(self):
        """初始化记忆管理器"""
        self._memories: Dict[str, BaseMemory] = {}
        self._memory_types: Dict[str, type] = {}
    
    def register_memory_type(self, memory_type: str, memory_class: type) -> None:
        """
        注册记忆类型
        
        Args:
            memory_type: 记忆类型名称
            memory_class: 记忆类
        """
        self._memory_types[memory_type] = memory_class
    
    def create_memory(self, agent_id: str, memory_type: str = "buffer",
                     llm: Optional[BaseLanguageModel] = None, **kwargs) -> BaseMemory:
        """
        创建记忆实例
        
        Args:
            agent_id: 智能体ID
            memory_type: 记忆类型
            llm: 大语言模型
            **kwargs: 其他参数
            
        Returns:
            记忆实例
        """
        if memory_type not in self._memory_types:
            raise ValueError(f"未知的记忆类型: {memory_type}")
        
        memory_class = self._memory_types[memory_type]
        memory = memory_class(agent_id, llm, **kwargs)
        
        self._memories[agent_id] = memory
        return memory
    
    def get_memory(self, agent_id: str) -> Optional[BaseMemory]:
        """
        获取记忆实例
        
        Args:
            agent_id: 智能体ID
            
        Returns:
            记忆实例或None
        """
        return self._memories.get(agent_id)
    
    def remove_memory(self, agent_id: str) -> bool:
        """
        移除记忆实例
        
        Args:
            agent_id: 智能体ID
            
        Returns:
            是否移除成功
        """
        if agent_id in self._memories:
            del self._memories[agent_id]
            return True
        return False
    
    def list_memories(self) -> List[str]:
        """
        列出所有记忆ID
        
        Returns:
            记忆ID列表
        """
        return list(self._memories.keys())
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """
        获取所有记忆的统计信息
        
        Returns:
            统计信息字典
        """
        stats = {
            "total_memories": len(self._memories),
            "memory_types": list(self._memory_types.keys()),
            "memories": {}
        }
        
        for agent_id, memory in self._memories.items():
            stats["memories"][agent_id] = memory.get_memory_stats()
        
        return stats
    
    def clear_all_memories(self) -> int:
        """
        清空所有记忆
        
        Returns:
            清空的记忆数量
        """
        count = len(self._memories)
        for memory in self._memories.values():
            memory.clear_memory()
        self._memories.clear()
        return count


# 全局记忆管理器
_memory_manager = MemoryManager()


def get_memory_manager() -> MemoryManager:
    """获取全局记忆管理器"""
    return _memory_manager


def create_memory(agent_id: str, memory_type: str = "buffer", 
                 llm: Optional[BaseLanguageModel] = None, **kwargs) -> BaseMemory:
    """创建记忆的便捷函数"""
    return _memory_manager.create_memory(agent_id, memory_type, llm, **kwargs)


def get_memory(agent_id: str) -> Optional[BaseMemory]:
    """获取记忆的便捷函数"""
    return _memory_manager.get_memory(agent_id)
